# app.py
from datetime import datetime
from flask_sqlalchemy import SQLAlchemy
from werkzeug.security import generate_password_hash, check_password_hash
from flask_login import UserMixin
from flask_sqlalchemy import SQLAlchemy
import uuid

db = SQLAlchemy()  # 创建数据库实例


# 模型定义
class User(UserMixin, db.Model):
    __tablename__ = 'users'
    user_id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    username = db.Column(db.String(80), unique=True, nullable=False)
    password_hash = db.Column(db.String(120), nullable=False)
    is_admin = db.Column(db.Boolean, default=False)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    last_login = db.Column(db.DateTime)

    agents = db.relationship('Agent', backref='creator', lazy=True)

    def set_password(self, password):
        self.password = generate_password_hash(password)

    def check_password(self, password):
        return check_password_hash(self.password, password)


class LLMModel(db.Model):
    __tablename__ = 'llm_models'
    model_id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    name = db.Column(db.String(120), nullable=False)
    model_type = db.Column(db.String(20), nullable=False)  # 'chat' or 'embedding'
    api_endpoint = db.Column(db.String(255), nullable=False)
    api_key = db.Column(db.String(120))
    max_tokens = db.Column(db.Integer, default=8192)
    supports_functions = db.Column(db.Boolean, default=False)
    is_active = db.Column(db.Boolean, default=True)
    description = db.Column(db.Text)
    provider = db.Column(db.String(80))

    agents = db.relationship('Agent', backref='llm_model', lazy=True)
    knowledge_bases = db.relationship('KnowledgeBase', backref='embedding_model', lazy=True)


class KnowledgeBase(db.Model):
    __tablename__ = 'knowledge_bases'
    kb_id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    name = db.Column(db.String(120), nullable=False)
    embedding_model_id = db.Column(db.String(36), db.ForeignKey('llm_models.model_id'), nullable=False)
    vector_db_path = db.Column(db.String(255), nullable=False)
    description = db.Column(db.Text)
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    creator_id = db.Column(db.String(36), db.ForeignKey('users.user_id'))
    documents_count = db.Column(db.Integer, default=0)
    size = db.Column(db.Integer, default=0)
    is_public = db.Column(db.Boolean, default=False)
    chunk_size = db.Column(db.Integer, default=512)
    chunk_overlap = db.Column(db.Integer, default=50)

    # # 新增字段
    # kid = db.Column(db.String(36), db.ForeignKey('knowledge_bases.kb_id'), nullable=False)

    agents = db.relationship('Agent', secondary='agent_knowledgebase', back_populates='knowledge_bases')


class Agent(db.Model):
    __tablename__ = 'agents'
    agent_id = db.Column(db.String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    name = db.Column(db.String(120), nullable=False)
    llm_model_id = db.Column(db.String(36), db.ForeignKey('llm_models.model_id'), nullable=False)
    creator_id = db.Column(db.String(36), db.ForeignKey('users.user_id'), nullable=False)
    greeting_message = db.Column(db.String(255), default="你好，我是AI助理，有什么可以帮您？")
    created_at = db.Column(db.DateTime, default=datetime.utcnow)
    updated_at = db.Column(db.DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
    is_active = db.Column(db.Boolean, default=True)
    system_prompt = db.Column(db.Text)
    temperature = db.Column(db.Float, default=0.7)
    max_tokens = db.Column(db.Integer, default=2048)

    knowledge_bases = db.relationship('KnowledgeBase', secondary='agent_knowledgebase', back_populates='agents')


class DocumentMetadata(db.Model):
    __tablename__ = 'document_metadata'

    id = db.Column(db.Integer, primary_key=True)
    doc_id = db.Column(db.String, unique=True)  # 与LlamaIndex的doc_id对应
    file_name = db.Column(db.String)
    file_type = db.Column(db.String)
    file_size = db.Column(db.Integer)
    page_count = db.Column(db.Integer)
    created_at = db.Column(db.String)
    additional_metadata = db.Column(db.JSON)  # 存储其他元数据
    summary = db.Column(db.Text)  # 文档摘要
    text_preview = db.Column(db.Text)  # 文本预览

    kid = db.Column(db.String(36), db.ForeignKey('knowledge_bases.kb_id', name='fk_document_metadata_kid'),
                    nullable=False)
    # kid = db.Column(db.String(36), db.ForeignKey('knowledge_bases.kb_id'), nullable=False)

# 关联表
agent_knowledgebase = db.Table('agent_knowledgebase',
                               db.Column('agent_id', db.String(36), db.ForeignKey('agents.agent_id'), primary_key=True),
                               db.Column('kb_id', db.String(36), db.ForeignKey('knowledge_bases.kb_id'),
                                         primary_key=True)
                               )
